{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " OpenAI Gym example using PySNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "import torch\n",
    "from ray import tune\n",
    "from ray.tune.schedulers import ASHAScheduler\n",
    "from ray.tune.suggest.bayesopt import BayesOptSearch\n",
    "\n",
    "from pysnn.network import SNNNetwork\n",
    "from pysnn.neuron import AdaptiveLIFNeuron, LIFNeuron\n",
    "from pysnn.connection import Linear\n",
    "from pysnn.learning import MSTDPET"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## Network definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SpikingNet(SNNNetwork):\n",
    "    def __init__(self, n_in_dynamics, n_hid_dynamics, n_out_dynamics, c_dynamics):\n",
    "        super(SpikingNet, self).__init__()\n",
    "\n",
    "        # Input layer\n",
    "        # 4-dimensional observation of + and -, so 8 neurons\n",
    "        self.neuron0 = AdaptiveLIFNeuron((1, 1, 8), *n_in_dynamics)\n",
    "\n",
    "        # Hidden layer\n",
    "        # Adaptive neuron to cope with highly varying input\n",
    "        self.neuron1 = AdaptiveLIFNeuron((1, 1, 64), *n_hid_dynamics)\n",
    "\n",
    "        # Output layer\n",
    "        # Non-adaptive neuron to decrease latency in action selection\n",
    "        # Action is binary, so single neuron suffices\n",
    "        self.neuron2 = LIFNeuron((1, 1, 1), *n_out_dynamics)\n",
    "\n",
    "        # Connections\n",
    "        # Weights initialized uniformly in [0, 1]\n",
    "        self.linear1 = Linear(8, 64, *c_dynamics)\n",
    "        self.linear2 = Linear(64, 1, *c_dynamics)\n",
    "\n",
    "        # Layers (for learning rule)\n",
    "        # Consist of connection and post-synaptic neuron\n",
    "        self.add_layer(\"fc1\", self.linear1, self.neuron1)\n",
    "        self.add_layer(\"fc2\", self.linear2, self.neuron2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Encode before passing to input layer\n",
    "        x = self._encode(x)\n",
    "        spikes, trace = self.neuron0(x)\n",
    "\n",
    "        # Hidden layer\n",
    "        # Connection trace (2nd argument) is not used\n",
    "        x, _ = self.linear1(spikes, trace)\n",
    "        spikes, trace = self.neuron1(x)\n",
    "\n",
    "        # Output layer\n",
    "        x, _ = self.linear2(spikes, trace)\n",
    "        spikes, trace = self.neuron2(x)\n",
    "\n",
    "        # Decode into binary action (so nothing)\n",
    "        return self._decode(spikes)\n",
    "\n",
    "    def _encode(self, x):\n",
    "        # Repeat the input\n",
    "        x = x.repeat(1, 1, 2)\n",
    "\n",
    "        # Clamp first half to positive, second to negative\n",
    "        x[..., :4].clamp_(min=0)\n",
    "        x[..., 4:].clamp_(max=0)\n",
    "\n",
    "        # Make absolute\n",
    "        return x.abs().float()\n",
    "\n",
    "    def _decode(self, x):\n",
    "        return x.byte()\n",
    "\n",
    "    def step(self, obs, env, rule, render=False):\n",
    "        # Observation is first positional argument\n",
    "        # Environment is second\n",
    "        # Learning rule is third\n",
    "        obs = torch.from_numpy(obs).view(1, 1, -1)\n",
    "        action = self.forward(obs)\n",
    "\n",
    "        # Optional render of environment\n",
    "        if render:\n",
    "            env.render()\n",
    "\n",
    "        # Do environment step\n",
    "        action = action.item()\n",
    "        obs, reward, done, _ = env.step(action)\n",
    "\n",
    "        # Do learning step\n",
    "        rule.update_state()\n",
    "        rule.step(reward)\n",
    "\n",
    "        # Return stepped environment and its returns\n",
    "        return obs, reward, done, env, rule"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## Training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main(config):\n",
    "    # Put config in neuron/connection/rule dicts\n",
    "    neuron_in = [\n",
    "        config[\"thresh0\"],\n",
    "        config[\"v_rest\"],\n",
    "        config[\"alpha_v0\"],\n",
    "        config[\"alpha_t0\"],\n",
    "        config[\"dt\"],\n",
    "        config[\"refrac\"],\n",
    "        config[\"tau_v0\"],\n",
    "        config[\"tau_t0\"],\n",
    "        config[\"alpha_thresh0\"],\n",
    "        config[\"tau_thresh0\"],\n",
    "    ]\n",
    "    neuron_hid = [\n",
    "        config[\"thresh1\"],\n",
    "        config[\"v_rest\"],\n",
    "        config[\"alpha_v1\"],\n",
    "        config[\"alpha_t1\"],\n",
    "        config[\"dt\"],\n",
    "        config[\"refrac\"],\n",
    "        config[\"tau_v1\"],\n",
    "        config[\"tau_t1\"],\n",
    "        config[\"alpha_thresh1\"],\n",
    "        config[\"tau_thresh1\"],\n",
    "    ]\n",
    "    neuron_out = [\n",
    "        config[\"thresh2\"],\n",
    "        config[\"v_rest\"],\n",
    "        config[\"alpha_v2\"],\n",
    "        config[\"alpha_t2\"],\n",
    "        config[\"dt\"],\n",
    "        config[\"refrac\"],\n",
    "        config[\"tau_v2\"],\n",
    "        config[\"tau_t2\"],\n",
    "    ]\n",
    "    conns = [config[\"batch_size\"], config[\"dt\"], config[\"delay\"]]\n",
    "    lr = [config[\"a_pre\"], config[\"a_post\"], config[\"lr\"], config[\"tau_e_trace\"]]\n",
    "\n",
    "    # Build network\n",
    "    # Build learning rule from network layers\n",
    "    network = SpikingNet(neuron_in, neuron_hid, neuron_out, conns)\n",
    "    rule = MSTDPET(network.layer_state_dict(), *lr)\n",
    "\n",
    "    # Build env\n",
    "    env = gym.make(\"CartPole-v1\")\n",
    "    obs = env.reset()\n",
    "\n",
    "    # Logging variables\n",
    "    episode_reward = 0.0\n",
    "\n",
    "    # Simulation loop\n",
    "    for step in range(config[\"steps\"]):\n",
    "        obs, reward, done, env, rule = network.step(obs, env, rule, render=False)\n",
    "        episode_reward += reward\n",
    "\n",
    "        # Episode end\n",
    "        if done:\n",
    "            obs = env.reset()\n",
    "            network.reset_state()\n",
    "            rule.reset_state()\n",
    "            tune.track.log(reward=episode_reward)\n",
    "            episode_reward = 0.0\n",
    "\n",
    "    # Cleanup\n",
    "    env.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## Execute training with hyperparameter optimization\n",
    " Makes use the bayesian-optimization package, `pip install bayesian-optimization` and ray[tune], `pip install ray[tune]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fixed parameters\n",
    "config = {\n",
    "    \"dt\": 1.0,\n",
    "    \"thresh0\": 0.2,\n",
    "    \"thresh1\": 0.2,\n",
    "    \"refrac\": 0,\n",
    "    \"v_rest\": 0.0,\n",
    "    \"batch_size\": 1,\n",
    "    \"delay\": 0,\n",
    "    \"a_post\": 0.0,\n",
    "    \"steps\": 10000,\n",
    "}\n",
    "\n",
    "# Search space for Bayesian Optimization\n",
    "space = {\n",
    "    \"thresh2\": (0.0, 1.0),\n",
    "    \"alpha_v0\": (0.0, 2.0),\n",
    "    \"alpha_v1\": (0.0, 2.0),\n",
    "    \"alpha_v2\": (0.0, 2.0),\n",
    "    \"alpha_t0\": (0.0, 2.0),\n",
    "    \"alpha_t1\": (0.0, 2.0),\n",
    "    \"alpha_t2\": (0.0, 2.0),\n",
    "    \"alpha_thresh0\": (0.0, 2.0),\n",
    "    \"alpha_thresh1\": (0.0, 2.0),\n",
    "    \"tau_v0\": (0.0, 1.0),\n",
    "    \"tau_v1\": (0.0, 1.0),\n",
    "    \"tau_v2\": (0.0, 1.0),\n",
    "    \"tau_t0\": (0.0, 1.0),\n",
    "    \"tau_t1\": (0.0, 1.0),\n",
    "    \"tau_t2\": (0.0, 1.0),\n",
    "    \"tau_thresh0\": (0.0, 1.0),\n",
    "    \"tau_thresh1\": (0.0, 1.0),\n",
    "    \"a_pre\": (0.0, 2.0),\n",
    "    \"lr\": (1e-6, 1e-2),\n",
    "    \"tau_e_trace\": (0.0, 1.0),\n",
    "}\n",
    "\n",
    "# Run hyperparameter search\n",
    "search = BayesOptSearch(\n",
    "    space,\n",
    "    max_concurrent=6,\n",
    "    metric=\"reward\",\n",
    "    mode=\"max\",\n",
    "    utility_kwargs={\"kind\": \"ucb\", \"kappa\": 2.5, \"xi\": 0.0},\n",
    ")\n",
    "scheduler = ASHAScheduler(metric=\"reward\", mode=\"max\")\n",
    "tune.run(\n",
    "    main,\n",
    "    num_samples=100,\n",
    "    scheduler=scheduler,\n",
    "    search_alg=search,\n",
    "    config=config,\n",
    "    verbose=1,\n",
    "    local_dir=\"ray_runs\",\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "file_extension": ".py",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
